#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
batch_asteroid_calmag_with_util.py

批量检测小行星，提取 CALMAG_12，绘制光变曲线，
生成小行星切片马赛克与中心切片 GIF。
"""

import os
import sys
import argparse
from datetime import datetime
from loguru import logger

# 天文工具库导入
from astroquery.imcce import Skybot
from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy.time import Time
import astropy.units as u
from astropy.table import Table
from astropy.wcs import WCS
from astropy.visualization import simple_norm

# 数学绘图模块
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.patches import Rectangle

# 工具函数导入
from tools.utils import TargetXCatalogue    # 交叉匹配函数
from tools.make_gif import make_gif         # GIF 生成函数
import pandas as pd
import matplotlib.ticker as ticker
import tqdm
ticker.MAXTICKS = 100000  # 防止时间刻度过多报错

# 从 CSV 文件加载小行星真值表
def load_known_asteroids_table(csv_path):
    df = pd.read_csv(csv_path, header=None)
    df.columns = [
        'filename', 'obs_time', 'id', 'x', 'y', 'magerr', 'mjd', 'exptime', 'mag',
        'radec_hms', 'true_ra', 'true_dec'
    ]
    df['obs_time'] = pd.to_datetime(df['obs_time'])
    return df

# 从本地真值表匹配或调用 SkyBot 获取小行星
def find_or_query_asteroid(fits_path, known_asteroids_df, radius_arcmin=80.0, time_tol_sec=30):
    hdr = fits.getheader(fits_path, 0)
    ra_hdr = hdr.get('CRVAL1')
    dec_hdr = hdr.get('CRVAL2')
    date = hdr.get('DATE-OBS')
    filename = os.path.basename(fits_path)

    if ra_hdr is None or dec_hdr is None or date is None:
        raise KeyError("FITS header 缺少 CRVAL1 / CRVAL2 / DATE-OBS")

    obs_time = pd.to_datetime(date)

    matched = known_asteroids_df[
        (known_asteroids_df['filename'] == filename) &
        (abs((known_asteroids_df['obs_time'] - obs_time).dt.total_seconds()) < time_tol_sec)
    ]
    if not matched.empty:
        return obs_time.to_pydatetime(), matched

    logger.error(f"⚠️ 未找到匹配记录，转用 SkyBot 查询, {fits_path}")
    coord = SkyCoord(float(ra_hdr) * u.deg, float(dec_hdr) * u.deg)
    epoch = Time(obs_time)
    result = Skybot.cone_search(coord, radius_arcmin * u.arcmin, epoch, 'O44')
    return obs_time.to_pydatetime(), result

# 加载 phot_tbl 并检查必要字段
def load_calmag_catalog(cat_path):
    """
    加载带有 CALMAG_12 的星表文件。
    自动检测第一个包含所需字段的 HDU。
    """
    from astropy.io import fits

    with fits.open(cat_path) as hdul:
        for hdu in hdul:
            if isinstance(hdu, fits.BinTableHDU) or isinstance(hdu, fits.TableHDU):
                try:
                    tbl = Table(hdu.data)
                    if all(k in tbl.colnames for k in ('ALPHA_J2000','DELTA_J2000','CALMAG_12','MAGERR_AUTO_S')):
                        return tbl
                except Exception as e:
                    continue  # 某些 HDU 可能不能转换为 Table
    raise ValueError(f"{cat_path} 中未找到包含必要字段的表格 HDU")

# 对小行星列表与 phot_tbl 进行交叉匹配
def crossmatch_with_util(ast_tbl, phot_tbl, max_sep_arcsec=2.0):
    results = []
    if isinstance(ast_tbl, pd.DataFrame):
        for _, row in ast_tbl.iterrows():
            ra, dec = float(row['true_ra']), float(row['true_dec'])
            match = TargetXCatalogue(ra, dec, phot_tbl, radius=max_sep_arcsec,RaKeyword='ALPHA_J2000', DecKeyword='DELTA_J2000')
            if match is not None:
                calmag = match['CALMAG_12']
                magerr = match['MAGERR_AUTO_S']
                flags = match['FLAGS']
            else:
                calmag = magerr = flags = None
            results.append((calmag, magerr, flags))
    else:
        for row in ast_tbl:
            ra, dec = row['RA'].value, row['DEC'].value
            match = TargetXCatalogue(ra, dec, phot_tbl, radius=max_sep_arcsec, RaKeyword='ALPHA_J2000', DecKeyword='DELTA_J2000')
            if match is not None:
                calmag = match['CALMAG_12']
                magerr = match['MAGERR_AUTO_S']
                flags = match['FLAGS']
            else:
                calmag = magerr = flags = None
            results.append((calmag, magerr,flags))
    return results

# 调试：可视化小行星和星表交叉点位置
def visualize_asteroid_position(fits_path, phot_tbl, ra, dec, save_path):
    img = fits.getdata(fits_path, 0)
    hdr = fits.getheader(fits_path, 0)
    wcs = WCS(hdr)

    fig, ax = plt.subplots(figsize=(8, 8), subplot_kw={'projection': wcs})
    norm = simple_norm(img, 'sqrt', percent=99)
    ax.imshow(img, origin='lower', cmap='gray', norm=norm)
    ax.scatter(phot_tbl['ALPHA_J2000'], phot_tbl['DELTA_J2000'], transform=ax.get_transform('world'), s=10, c='cyan', label='Catalog sources')
    ax.scatter([ra], [dec], transform=ax.get_transform('world'), s=50, c='red', marker='x', label='Asteroid')
    ax.legend()
    ax.set_title('Asteroid Position vs Catalog')
    plt.savefig(save_path, dpi=200)
    plt.close()

# 主流程：遍历所有 fits 图像，处理 photometry 与交叉匹配
def process_directory(root_dir, radius=22.0, max_sep=2.0):
    all_data = []
    all_data_with_mag = []
    for dirpath, _, files in os.walk(root_dir):
        for fn in files:
            if not fn.lower().endswith('sciimg.fits'):
                continue
            path = os.path.join(dirpath, fn)
            band = fits.getheader(path)["FILTER"]
            try:
                df_known = load_known_asteroids_table(f'/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/images50cm/asteroid/12377_{band}.csv')  # 根据 band 加载真值表
                obs_time, ast_tbl = find_or_query_asteroid(path, df_known)
            except Exception as e:
                logger.error(f"{fn}: 小行星检测失败：{e}")
                continue
            base = fn[:-5]
            cat_path = os.path.join(dirpath, f"{base}_sexcat_CalMag.fits")
            if not os.path.exists(cat_path):
                continue
            phot_tbl = load_calmag_catalog(cat_path)
            matches = crossmatch_with_util(ast_tbl, phot_tbl, max_sep)
            for idx, (calmag, magerr,flags) in enumerate(matches):
                if calmag is not None:
                    all_data.append((obs_time, band, calmag, magerr))
                # 可视化当前小行星与星表对比
                if isinstance(ast_tbl, pd.DataFrame):
                    ra, dec = float(ast_tbl.iloc[idx]['true_ra']), float(ast_tbl.iloc[idx]['true_dec'])
                else:
                    ra, dec = ast_tbl[idx]['RA'].value, ast_tbl[idx]['DEC'].value
                # vis_path = os.path.join("result50cm","vis", f"vis_{fn[:-5]}_{idx}.png")
                # visualize_asteroid_position(path, phot_tbl, ra, dec, vis_path)
                all_data_with_mag.append((fn,obs_time, band, calmag, magerr,ra, dec,flags))
    return all_data,all_data_with_mag

# 绘制光变曲线图，按日期为每晚生成一个 subplot

def plot_time_mag(data, separate=False):
    from matplotlib.dates import DateFormatter, AutoDateLocator
    from collections import defaultdict

    by_day = defaultdict(list)
    for d in data:
        key = d[0].date()
        by_day[key].append(d)

    nrow = len(by_day)
    fig, axes = plt.subplots(nrow, 1, figsize=(12, 3*nrow), sharex=False)
    if nrow == 1:
        axes = [axes]

    for ax, (day, records) in zip(axes, sorted(by_day.items())):
        bands = sorted(set(r[1] for r in records))
        for band in bands:
            times = [r[0] for r in records if r[1] == band]
            mags = [r[2] for r in records if r[1] == band]
            errs = [r[3] for r in records if r[1] == band]
            ax.errorbar(times, mags, yerr=errs, fmt='o', capsize=3, label=band)
        ax.set_title(f"Asteroid Light Curve on {day}")
        ax.invert_yaxis()
        ax.set_ylabel("CALMAG_12")
        ax.grid(True, which='both', axis='x', linestyle='--', alpha=0.5)
        locator = AutoDateLocator(maxticks=10)
        ax.xaxis.set_major_locator(locator)
        ax.xaxis.set_major_formatter(DateFormatter('%H:%M'))
        ax.legend(title='Band')

    axes[-1].set_xlabel("Observation Time")
    plt.tight_layout()
    plt.savefig("result50cm/time_mag_by_day.png", dpi=300)
    plt.close(fig)

def make_asteroid_mosaic(root_dir, output_path, size=200, max_per_row=5,radius=22.0, max_sep=2.0, band_order=('g', 'r')):
    """
    遍历 root_dir 下所有 sciimg.fits，
    对每个文件中小行星位置截取 size x size 区域，
    并按网格保存到 output_path。
    每个子图顶部显示：
        <filename>
        <band>-band: mag=XX.XX±YY.YY
    """
    entries = []
    for dirpath, _, files in tqdm.tqdm(os.walk(root_dir)):
        for fn in files:
            if not fn.lower().endswith('sciimg.fits'):
                continue
            fits_path = os.path.join(dirpath, fn)
            band = fits.getheader(fits_path)["FILTER"]
            try:
                df_known = load_known_asteroids_table(f'/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/images50cm/asteroid/12377_{band}.csv')  # 根据 band 加载真值表
                obs_time, ast_tbl = find_or_query_asteroid(fits_path, df_known)
                if isinstance(ast_tbl, pd.DataFrame):
                    row = ast_tbl.iloc[0]
                    ra, dec = float(row['true_ra']), float(row['true_dec'])
                else:
                    row = next(iter(ast_tbl))
                    ra, dec = row['RA'].value, row['DEC'].value
                base = fn[:-5]
                cat_path = os.path.join(dirpath, f"{base}_sexcat_CalMag.fits")
                if not os.path.exists(cat_path):
                    continue
                phot_tbl = load_calmag_catalog(cat_path)
                calmag, magerr, flags= crossmatch_with_util(ast_tbl, phot_tbl, max_sep)[0]

                hdul = fits.open(fits_path)
                data = hdul[0].data
                wcs = WCS(hdul[0].header)
                xpix, ypix = wcs.world_to_pixel(SkyCoord(ra*u.deg, dec*u.deg))
                x, y = int(xpix), int(ypix)
                half = size // 2
                cut = data[y-half:y+half, x-half:x+half]

                if calmag is None or magerr is None:
                    label = f"{fn}\n{band}-band: mag=N/A"
                else:
                    label = f"{fn}\n{band}-band: {calmag:.2f}±{magerr:.2f}"
                entries.append((band, obs_time, label, cut))
            except Exception as e:
                logger.warning(f"跳过 {fn}: {e}")
                continue

    if not entries:
        logger.warning("未找到任何裁剪图用于拼图")
        return

    entries.sort(key=lambda e: (band_order.index(e[0]) if e[0] in band_order else len(band_order), e[1]))
    entries=entries[:100]
    n = len(entries)
    ncol = min(max_per_row, n)
    nrow = int(np.ceil(n / ncol))
    fig, axes = plt.subplots(nrow, ncol, figsize=(2.5*ncol, 2.5*nrow))
    axes = np.atleast_1d(axes).flatten()
    for ax in axes[n:]:
        ax.axis('off')
    for i, (_, _, label, cut) in enumerate(entries):
        ax = axes[i]
        norm = simple_norm(cut, 'sqrt', percent=99)
        ax.imshow(cut, norm=norm, cmap='gray', origin='lower')
        ax.set_title(label, fontsize=6)
        ax.axis('off')
    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.close(fig)

def make_region_cutout(fits_path, x_pix, y_pix, output_dir, size=500,filiename=None):
    """
    从单帧 FITS 中按给定像素位置裁剪 size×size 区域，
    并在中心位置画红框标记小行星，输出 PNG。
    """
    data = fits.getdata(fits_path)
    
    time_ser = fits.getheader(fits_path)["DATE-OBS"].replace("-","").replace(" ","").replace(":","")
    half = size // 2
    cut  = data[y_pix-half:y_pix+half, x_pix-half:x_pix+half]
    norm = simple_norm(cut, 'sqrt', percent=99)

    fig, ax = plt.subplots(figsize=(5,5), dpi=100)
    ax.imshow(cut, norm=norm, cmap='gray', origin='lower')

    # 在中心位置绘制 20×20 的红框
    rect = Rectangle((half-10, half-10), 20, 20,edgecolor='red', fill=False, lw=1.5)
    ax.add_patch(rect)
    ax.axis('off')

    os.makedirs(output_dir, exist_ok=True)
    outfn = time_ser+"_"+os.path.splitext(os.path.basename(fits_path))[0] + '.png'
    fig.savefig(os.path.join(output_dir, outfn),
                bbox_inches='tight', pad_inches=0)
    plt.close(fig)
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="批量检测小行星、提取 CALMAG、绘制光变曲线、生成切片与马赛克/GIF"
    )
    parser.add_argument("-d", "--directory",
                        default="./images50cm/12377/",
                        help="根目录，遍历 sciimg.fits")
    parser.add_argument("-r", "--radius", type=float, default=22.0,
                        help="Skybot 搜索半径（arcmin）")
    parser.add_argument("-s", "--maxsep", type=float, default=2.0,
                        help="交叉匹配最大分离（arcsec）")
    parser.add_argument("--separate", default="True",
                        help="为各波段单独子图")
    parser.add_argument("--mosaic", default=False,
                        help="生成小行星切片马赛克")
    parser.add_argument("--cutout", default=True,
                        help="生成每帧中心切片并标记")
    parser.add_argument("--size", type=int, default=200,
                        help="切片/马赛克单元大小（px）")
    parser.add_argument("--maxcol", type=int, default=5,
                        help="马赛克最大列数")
    parser.add_argument("--fps", type=float, default=2.0,
                        help="中心切片 GIF 帧率 (frame/sec)")
    args = parser.parse_args()

    
    # 依赖检查
    try:
        import astroquery  # noqa
        import astropy     # noqa
        import matplotlib  # noqa
    except ImportError as ie:
        sys.exit(f"缺少依赖：{ie.name}，请安装 astroquery astropy matplotlib")

    os.makedirs("result50cm", exist_ok=True)

    # 1. 绘制光变曲线
    data,all_data_with_radec = process_directory(args.directory, args.radius, args.maxsep)
    if data:
        logger.info(f"data:{len(data)}")
        plot_time_mag(data, args.separate)
        logger.info("光变曲线图输出到 result50cm/")
        # 转换为 DataFrame
        df = pd.DataFrame(all_data_with_radec, columns=["fn", "obs_time", "band", "calmag", "magerr","ra","dec","flags"])
        # 保存为 CSV
        df.to_csv("result50cm/asteroid_photometry_results.csv", index=False)

    # 2. 生成马赛克
    if args.mosaic:
        for band_ in ["g","r"]:
            outm = f"result50cm/mosaic_{band_}.png"
            make_asteroid_mosaic(args.directory, outm,size=args.size, max_per_row=args.maxcol,band_order=(band_,))
            logger.info(f"马赛克图输出到 {outm}")

    # 3. 生成中心切片并合成 GIF
    if args.cutout:
        for dirpath, _, files in tqdm.tqdm(os.walk(args.directory)):
            for fn in files:
                try:
                    if not fn.lower().endswith("sciimg.fits"):
                        continue
                    path = os.path.join(dirpath, fn)
                    band=fits.getheader(path)["FILTER"]
                    df_known = load_known_asteroids_table(f'/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/images50cm/asteroid/12377_{band}.csv')
                    fits_file = path
                    obs_time, ast_tbl = find_or_query_asteroid(fits_file, df_known)
                    
                    if isinstance(ast_tbl, pd.DataFrame):
                        if ast_tbl is None or len(ast_tbl) == 0:
                            continue
                        ra, dec = ast_tbl['true_ra'], ast_tbl['true_dec']
                    else:
                        if not ast_tbl:
                            continue
                        ra, dec = ast_tbl[0]['RA'].value, ast_tbl[0]['DEC'].value

                    # 对应 FITS 中小行星第一颗的位置
                    
                    wcs = WCS(fits.getheader(path))
                    xpix, ypix = wcs.world_to_pixel(SkyCoord(ra*u.deg, dec*u.deg))

                    # 按滤波器分目录保存
                    filt  = fits.getheader(path)["FILTER"]
                    cutdir = f"result50cm/cutouts/{filt}"
                    make_region_cutout(path, int(xpix), int(ypix), cutdir, size=500)
                except:
                    pass

        # 生成 GIF：请带上后缀 .gif
        for filt in ["g", "r"]:
            cutdir = f"result50cm/cutouts/{filt}"
            if os.path.join(cutdir):
                gif_path = f"result50cm/asteroid_{filt}.gif"
                make_gif(cutdir, gif_path, pattern="*.png", fps=args.fps)
                logger.info(f"中心切片 GIF 输出到 {gif_path}")